from asyncio import Future
from typing import Any, Coroutine

from paddleocr import PaddleOCR
import asyncio
from . import utils
from . import objectPool
from .config import ThreadPool, logger, objectPoolSize, det, rec


def create_new_object() -> PaddleOCR:
    return PaddleOCR(use_angle_cls=True, lang="ch", page_num=1,
                     det=det, rec=rec, use_tensorrt=True, enable_mkldnn=False,
                     # cpu_threads=10,
                     # use_npu=True,
                     show_log=True)


class PaddleOCRUtil(metaclass=utils.Singleton):

    def __init__(self):
        self.paddleOcrPoll = objectPool.ObjectPool(create_new_object, objectPoolSize)

    # @utils.calc_self_time
    async def image_ocr(self, img_path) -> tuple[Any, list]:
        _loop = asyncio.get_running_loop()
        data: list = list()
        logger.debug(f"开始解析{img_path}")
        logger.debug(f"对象池{self.paddleOcrPoll.size()}")
        ocr = await self.paddleOcrPoll.acquire()
        try:
            result = await _loop.run_in_executor(ThreadPool, lambda: ocr.ocr(img_path, cls=True))
            for idx in range(len(result)):
                res = result[idx]
                for line in res:
                    data.append(line)
        finally:
            # lock.release()
            self.paddleOcrPoll.release(ocr)
            logger.debug(f"解析{img_path}完成")
        return img_path, data


class PaddleOCRService(metaclass=utils.Singleton):

    def __init__(self) -> None:
        self.paddleOCRUtil = PaddleOCRUtil()

    async def parserImage_run(self, files) -> dict:
        _loop = asyncio.get_running_loop()
        tasks = [_loop.create_task(self.paddleOCRUtil.image_ocr(i)) for i in files]
        results = {}
        for image_ocr in await asyncio.gather(*tasks):
            strs = []
            filename, data = image_ocr
            for i in data:
                strs.append(i[1][0])
            results[filename.split("\\")[-1]] = strs
        return results
